import re
import warnings

import matplotlib
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import seaborn as sns
import statsmodels.formula.api as smf
from lifelines import KaplanMeierFitter
from statsmodels.duration.hazard_regression import PHReg
from statsmodels.tools.sm_exceptions import ConvergenceWarning

# 设置样式
plt.style.use('seaborn-v0_8-whitegrid')
sns.set_palette("Set2")

# 忽略警告
warnings.simplefilter('ignore', ConvergenceWarning)
matplotlib.use('Agg')  # 使用Agg后端，避免显示问题

# 设置中文字体显示
plt.rcParams['font.sans-serif'] = ['Microsoft YaHei', 'SimHei', 'Arial Unicode MS']
plt.rcParams['axes.unicode_minus'] = False


def parse_week(week_str):
    if isinstance(week_str, str):
        match = re.match(r"(\d+)w\+(\d+)", week_str)
        if match:
            weeks = int(match.group(1))
            days = int(match.group(2))
            return weeks + days / 7
    return np.nan


def prepare_data():
    """准备数据"""
    # 读取数据
    df_male = pd.read_excel('附件.xlsx', sheet_name='男胎检测数据')

    # 数据清洗
    df_male['孕周'] = df_male['检测孕周'].apply(parse_week)
    df_male['Y浓度'] = df_male['Y染色体浓度']

    # 过滤无效数据
    valid_data = df_male.dropna(subset=['孕周', 'Y浓度', '孕妇BMI', '孕妇代码'])
    valid_data = valid_data[(valid_data['孕周'] > 5) & (valid_data['孕周'] < 40)]

    # 重命名列
    valid_data = valid_data.rename(columns={
        '孕妇代码': 'subject_id',
        '孕妇BMI': 'bmi',
        '年龄': 'age',
        '身高': 'height',
        '体重': 'weight'
    })

    return valid_data


def create_survival_data(df_long):
    """创建生存分析数据"""
    survival_data = []
    unique_subjects = df_long['subject_id'].unique()

    for subject in unique_subjects:
        subject_df = df_long[df_long['subject_id'] == subject].sort_values('孕周')

        # 检查是否达到阈值
        reached_threshold = subject_df[subject_df['Y浓度'] >= 0.04]

        if not reached_threshold.empty:
            # 找到首次达标时间
            event_time = reached_threshold['孕周'].min()
            status = 1  # 事件发生

            # 找到达标前最后一次检测时间
            before_event = subject_df[subject_df['孕周'] < event_time]
            prev_week = before_event['孕周'].max() if not before_event.empty else np.nan
        else:
            # 右删失情况
            event_time = subject_df['孕周'].max()
            status = 0  # 删失
            prev_week = np.nan

        # 收集协变量
        bmi_val = subject_df['bmi'].mean()
        age_val = subject_df['age'].mean()
        height_val = subject_df['height'].mean()
        weight_val = subject_df['weight'].mean()

        survival_data.append({
            'subject_id': subject,
            'event_time': event_time,
            'status': status,
            'left_interval': prev_week,
            'bmi': bmi_val,
            'age': age_val,
            'height': height_val,
            'weight': weight_val
        })

    return pd.DataFrame(survival_data)


def fit_mixed_effects_model(df_long):
    print("正在拟合混合效应模型...")

    # 使用更多协变量
    model_formula = "Y浓度 ~ 孕周 + bmi + age + height + weight"

    # 尝试不同的优化方法
    try:
        model = smf.mixedlm(model_formula, df_long,
                            groups=df_long["subject_id"],
                            re_formula="~孕周")
        result = model.fit(method='bfgs', maxiter=1000)
    except:
        # 如果失败，尝试更简单的方法
        model = smf.mixedlm(model_formula, df_long,
                            groups=df_long["subject_id"])
        result = model.fit(method='nm', maxiter=1000)

    print("混合效应模型拟合完成")
    print(result.summary())

    return result


def fit_cox_model(df_surv, longitudinal_predictions):
    """拟合Cox比例风险模型"""
    print("正在拟合Cox比例风险模型...")

    # 准备协变量
    exog_vars = df_surv[['bmi', 'age', 'height', 'weight']].copy()
    exog_vars['pred_y'] = longitudinal_predictions

    # 确保没有缺失值
    complete_cases = exog_vars.notna().all(axis=1)
    exog_vars = exog_vars[complete_cases]
    event_times = df_surv['event_time'][complete_cases]
    status = df_surv['status'][complete_cases]

    # 拟合Cox模型
    ph_model = PHReg(endog=event_times,
                     exog=exog_vars,
                     status=status)
    ph_result = ph_model.fit()

    print("Cox比例风险模型拟合完成")
    print(ph_result.summary())

    return ph_result


def plot_survival_curves(df_surv):
    """绘制生存曲线"""

    # BMI分组
    bins = [20, 28, 32, 36, 40, np.inf]
    labels = ['20-28', '28-32', '32-36', '36-40', '40+']
    df_surv['bmi_group'] = pd.cut(df_surv['bmi'], bins=bins, labels=labels, include_lowest=True)

    # 创建图形
    plt.figure(figsize=(10, 6))

    # 为每个BMI组绘制Kaplan-Meier曲线
    for group in labels:
        group_df = df_surv[df_surv['bmi_group'] == group]
        if len(group_df) > 0:
            kmf = KaplanMeierFitter()
            kmf.fit(durations=group_df['event_time'],
                    event_observed=group_df['status'],
                    label=f'BMI {group}')
            kmf.plot_survival_function(ci_show=False)

    plt.title('按BMI分组的生存曲线', fontsize=14)
    plt.xlabel('孕周', fontsize=12)
    plt.ylabel('未达标概率', fontsize=12)
    plt.legend(title='BMI分组')
    plt.tight_layout()
    plt.savefig('question three 生存曲线1.png', dpi=300)
    plt.close()

    print("生存曲线已保存为 question three1 生存曲线.png")


def calculate_optimal_times(df_surv, threshold=0.9):
    """计算最佳检测时点"""
    print("正在计算最佳检测时点...")

    # BMI分组
    bins = [20, 28, 32, 36, 40, np.inf]
    labels = ['20-28', '28-32', '32-36', '36-40', '40+']
    df_surv['bmi_group'] = pd.cut(df_surv['bmi'], bins=bins, labels=labels, include_lowest=True)

    optimal_times = {}

    for group in labels:
        group_df = df_surv[df_surv['bmi_group'] == group]
        if len(group_df) > 0:
            kmf = KaplanMeierFitter()
            kmf.fit(durations=group_df['event_time'],event_observed=group_df['status'])

            # 找到达到阈值的时间点
            survival_func = kmf.survival_function_
            cumulative_incidence = 1 - survival_func.iloc[:, 0]

            # 只考虑孕周大于10周的情况
            valid_times = cumulative_incidence.index[cumulative_incidence.index > 10]
            if len(valid_times) > 0:
                cumulative_incidence_valid = cumulative_incidence.loc[valid_times]
                reached_condition = cumulative_incidence_valid >= threshold

                if reached_condition.any():
                    optimal_time = cumulative_incidence_valid[reached_condition].index.min()
                    optimal_times[group] = optimal_time
                else:
                    optimal_times[group] = np.nan
            else:
                optimal_times[group] = np.nan
        else:
            optimal_times[group] = np.nan

    # 打印结果
    print("各BMI组的最佳检测时点 (达标概率 ≥ {}):".format(threshold))
    for group, time in optimal_times.items():
        if not np.isnan(time):
            print(f"BMI {group}: {time:.2f} 周")
        else:
            print(f"BMI {group}: 未达到阈值")

    return optimal_times


def sensitivity_analysis(df_long, df_surv, sigma_list=[0.01, 0.05, 0.1]):
    """敏感性分析：评估测量误差对结果的影响"""
    print("正在进行敏感性分析...")

    results = []
    original_optimal_times = calculate_optimal_times(df_surv)

    for sigma in sigma_list:
        print(f"分析噪声水平: σ = {sigma}")

        # 添加噪声
        noisy_long = df_long.copy()
        noise = np.random.normal(0, sigma, len(noisy_long))
        noisy_long['Y浓度'] += noise

        # 拟合混合效应模型
        try:
            noisy_model = smf.mixedlm("Y浓度 ~ 孕周 + bmi + age + height + weight",
                                      noisy_long,
                                      groups=noisy_long["subject_id"])
            noisy_result = noisy_model.fit(method='nm', maxiter=1000)

            # 预测浓度
            noisy_long['pred_y'] = noisy_result.predict(noisy_long)

            # 计算每个个体的平均预测浓度
            mean_pred_y = noisy_long.groupby('subject_id')['pred_y'].mean()

            # 更新生存数据
            noisy_surv = df_surv.copy()
            noisy_surv = noisy_surv.merge(mean_pred_y.rename('mean_pred_y'),
                                          left_on='subject_id', right_index=True, how='left')

            # 计算最佳检测时点
            optimal_times = calculate_optimal_times(noisy_surv)
            results.append((sigma, optimal_times))

        except Exception as e:
            print(f"在噪声水平 σ = {sigma} 下拟合失败: {e}")
            results.append((sigma, None))

    # 打印敏感性分析结果
    print("\n敏感性分析结果:")
    for sigma, optimal_times in results:
        if optimal_times is not None:
            print(f"噪声水平 σ = {sigma}:")
            for group, time in optimal_times.items():
                if not np.isnan(time):
                    print(f"  BMI {group}: {time:.2f} 周")
                else:
                    print(f"  BMI {group}: 未达到阈值")
        else:
            print(f"噪声水平 σ = {sigma}: 分析失败")

    return results




# 准备数据
df_long = prepare_data()
print(f"数据准备完成，共有 {len(df_long)} 条记录，{df_long['subject_id'].nunique()} 个个体")

# 创建生存数据
df_surv = create_survival_data(df_long)
print(f"生存数据创建完成，共有 {len(df_surv)} 个个体")

# 第一阶段：拟合混合效应模型
mixed_result = fit_mixed_effects_model(df_long)

# 预测浓度
df_long['pred_y'] = mixed_result.predict(df_long)

# 计算每个个体的平均预测浓度
mean_pred_y = df_long.groupby('subject_id')['pred_y'].mean()
df_surv = df_surv.merge(mean_pred_y.rename('mean_pred_y'),
                        left_on='subject_id', right_index=True, how='left')

# 第二阶段：拟合Cox比例风险模型
cox_result = fit_cox_model(df_surv, df_surv['mean_pred_y'])

# 绘制生存曲线
plot_survival_curves(df_surv)

# 计算最佳检测时点
optimal_times = calculate_optimal_times(df_surv)

# 敏感性分析
sensitivity_results = sensitivity_analysis(df_long, df_surv)



